"""
并行执行器类，用于管理多代理的并行处理
"""

import asyncio
import concurrent.futures
from typing import List, Dict, Any, Callable, Optional, Union
from dataclasses import dataclass
from loguru import logger
import time


@dataclass
class TaskResult:
    """任务结果"""
    task_id: str
    agent_name: str
    success: bool
    result: Any
    error: Optional[str] = None
    execution_time: float = 0.0


class ParallelExecutor:
    """并行执行器"""
    
    def __init__(self, max_workers: int = 4, timeout: int = 120):
        self.max_workers = max_workers
        self.timeout = timeout
        
    def execute_parallel(
        self,
        tasks: List[Dict[str, Any]],
        task_func: Callable,
        **kwargs
    ) -> List[TaskResult]:
        """并行执行任务列表"""
        results = []
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # 提交所有任务
            future_to_task = {}
            for i, task in enumerate(tasks):
                task_id = task.get('task_id', f'task_{i}')
                agent_name = task.get('agent_name', f'agent_{i}')
                
                future = executor.submit(
                    self._execute_single_task,
                    task_id,
                    agent_name,
                    task_func,
                    task,
                    **kwargs
                )
                future_to_task[future] = (task_id, agent_name)
            
            # 收集结果 - 使用更安全的方式处理超时
            completed_futures = set()
            try:
                for future in concurrent.futures.as_completed(future_to_task, timeout=self.timeout):
                    completed_futures.add(future)
                    task_id, agent_name = future_to_task[future]
                    try:
                        result = future.result()
                        results.append(result)
                        logger.info(f"任务 {task_id} ({agent_name}) 完成")
                    except Exception as e:
                        logger.error(f"任务 {task_id} ({agent_name}) 执行失败: {e}")
                        results.append(TaskResult(
                            task_id=task_id,
                            agent_name=agent_name,
                            success=False,
                            result=None,
                            error=str(e)
                        ))
            except concurrent.futures.TimeoutError:
                logger.warning(f"并行执行超时 ({self.timeout}s)，正在处理未完成的任务")
            
            # 处理未完成的futures
            unfinished_futures = set(future_to_task.keys()) - completed_futures
            if unfinished_futures:
                logger.warning(f"发现 {len(unfinished_futures)} 个未完成的任务，正在强制取消")
                
                for future in unfinished_futures:
                    task_id, agent_name = future_to_task[future]
                    
                    # 尝试取消未开始的任务
                    if future.cancel():
                        logger.info(f"成功取消任务 {task_id} ({agent_name})")
                        results.append(TaskResult(
                            task_id=task_id,
                            agent_name=agent_name,
                            success=False,
                            result=None,
                            error="任务被取消（超时）"
                        ))
                    else:
                        # 对于正在运行的任务，等待短时间后标记为超时
                        logger.warning(f"无法取消正在运行的任务 {task_id} ({agent_name})，等待结果")
                        try:
                            result = future.result(timeout=5)  # 额外等待5秒
                            results.append(result)
                            logger.info(f"任务 {task_id} ({agent_name}) 延迟完成")
                        except concurrent.futures.TimeoutError:
                            logger.error(f"任务 {task_id} ({agent_name}) 最终超时")
                            results.append(TaskResult(
                                task_id=task_id,
                                agent_name=agent_name,
                                success=False,
                                result=None,
                                error="任务执行超时"
                            ))
                        except Exception as e:
                            logger.error(f"任务 {task_id} ({agent_name}) 执行失败: {e}")
                            results.append(TaskResult(
                                task_id=task_id,
                                agent_name=agent_name,
                                success=False,
                                result=None,
                                error=str(e)
                            ))
        
        logger.info(f"并行执行完成，共处理 {len(results)} 个任务")
        return results
    
    def _execute_single_task(
        self,
        task_id: str,
        agent_name: str,
        task_func: Callable,
        task_data: Dict[str, Any],
        **kwargs
    ) -> TaskResult:
        """执行单个任务"""
        start_time = time.time()
        
        try:
            logger.info(f"开始执行任务 {task_id} ({agent_name})")
            result = task_func(task_data, **kwargs)
            execution_time = time.time() - start_time
            
            logger.info(f"任务 {task_id} ({agent_name}) 执行成功，耗时 {execution_time:.2f}s")
            return TaskResult(
                task_id=task_id,
                agent_name=agent_name,
                success=True,
                result=result,
                execution_time=execution_time
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            logger.error(f"任务 {task_id} ({agent_name}) 执行失败: {e}")
            return TaskResult(
                task_id=task_id,
                agent_name=agent_name,
                success=False,
                result=None,
                error=str(e),
                execution_time=execution_time
            )
    
    async def execute_async(
        self,
        tasks: List[Dict[str, Any]],
        task_func: Callable,
        **kwargs
    ) -> List[TaskResult]:
        """异步并行执行任务"""
        semaphore = asyncio.Semaphore(self.max_workers)
        
        async def execute_with_semaphore(task: Dict[str, Any]) -> TaskResult:
            async with semaphore:
                task_id = task.get('task_id', f'task_{id(task)}')
                agent_name = task.get('agent_name', f'agent_{id(task)}')
                
                return await self._execute_async_task(
                    task_id, agent_name, task_func, task, **kwargs
                )
        
        # 创建所有任务
        coroutines = [execute_with_semaphore(task) for task in tasks]
        
        # 并行执行
        try:
            results = await asyncio.wait_for(
                asyncio.gather(*coroutines, return_exceptions=True),
                timeout=self.timeout
            )
            
            # 处理异常结果
            processed_results = []
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    task = tasks[i]
                    task_id = task.get('task_id', f'task_{i}')
                    agent_name = task.get('agent_name', f'agent_{i}')
                    
                    processed_results.append(TaskResult(
                        task_id=task_id,
                        agent_name=agent_name,
                        success=False,
                        result=None,
                        error=str(result)
                    ))
                else:
                    processed_results.append(result)
            
            return processed_results
            
        except asyncio.TimeoutError:
            logger.error("异步任务执行超时")
            return [
                TaskResult(
                    task_id=task.get('task_id', f'task_{i}'),
                    agent_name=task.get('agent_name', f'agent_{i}'),
                    success=False,
                    result=None,
                    error="异步执行超时"
                )
                for i, task in enumerate(tasks)
            ]
    
    async def _execute_async_task(
        self,
        task_id: str,
        agent_name: str,
        task_func: Callable,
        task_data: Dict[str, Any],
        **kwargs
    ) -> TaskResult:
        """执行单个异步任务"""
        start_time = time.time()
        
        try:
            logger.info(f"开始执行异步任务 {task_id} ({agent_name})")
            
            # 如果task_func是协程函数，直接await
            if asyncio.iscoroutinefunction(task_func):
                result = await task_func(task_data, **kwargs)
            else:
                # 如果是同步函数，在线程池中执行
                loop = asyncio.get_event_loop()
                result = await loop.run_in_executor(
                    None, lambda: task_func(task_data, **kwargs)
                )
            
            execution_time = time.time() - start_time
            logger.info(f"异步任务 {task_id} ({agent_name}) 执行成功，耗时 {execution_time:.2f}s")
            
            return TaskResult(
                task_id=task_id,
                agent_name=agent_name,
                success=True,
                result=result,
                execution_time=execution_time
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            logger.error(f"异步任务 {task_id} ({agent_name}) 执行失败: {e}")
            return TaskResult(
                task_id=task_id,
                agent_name=agent_name,
                success=False,
                result=None,
                error=str(e),
                execution_time=execution_time
            )
    
    def get_execution_stats(self, results: List[TaskResult]) -> Dict[str, Any]:
        """获取执行统计信息"""
        total_tasks = len(results)
        successful_tasks = sum(1 for r in results if r.success)
        failed_tasks = total_tasks - successful_tasks
        
        execution_times = [r.execution_time for r in results if r.execution_time > 0]
        avg_execution_time = sum(execution_times) / len(execution_times) if execution_times else 0
        max_execution_time = max(execution_times) if execution_times else 0
        min_execution_time = min(execution_times) if execution_times else 0
        
        return {
            "total_tasks": total_tasks,
            "successful_tasks": successful_tasks,
            "failed_tasks": failed_tasks,
            "success_rate": successful_tasks / total_tasks if total_tasks > 0 else 0,
            "avg_execution_time": avg_execution_time,
            "max_execution_time": max_execution_time,
            "min_execution_time": min_execution_time,
            "total_execution_time": sum(execution_times)
        }